#!/bin/sh

# Peak Signal to Noise Ratio script for comparing the PSNR between two video files

# This script calculates the PSNR between two video files. The first video file is
# generally the original video, and the second is a modified version of the original.
# Often, the second video is encoded in a different codec, or uses filters to
# improve or change the video. The aim of this script is to give concrete numbers
# to often subjective video quality comparisons.
#
# The script calls psnrcore, which sequentially compares frames from both videos, 
# calculating the PSNR for each frame, and finally averages the overall PSNR for 
# both videos. Frame-by-frame data are written to a text file, while the final PSNR 
# is retured to standard out.

# INPUT:  (1) a pointer to the original video file.
#         (2) a pointer to the modified video file (to be compared against the 
#             original video file)
#         (3) a pointer to a file to which the PSNR for each frame will be written.
#         (4) the number of frames to compare, starting from the first frame.
#         (5) options to pass to MPlayer when creating frames from the original
#             video. Useful to add filters.
#         (6) options to pass to MPlayer when creating frames from the comparison
#             video. Ueeful to add filters.
#
# OUTPUT: (1) the overall, averaged PSNR between the two input directories.
#         (2) a log file of the PSNR for every frame (as given by the input pointer).

# Copyright (C) 2005 Joe Friedrichsen <pengi.films@gmail.com>
# Original shell script by Matthias Wieser, available in MPlayer CVS.
# Modified on 2005 September 23.
# 
# This program is free software; you can redistribute it and/or 
# modify it under the terms of the GNU General Public License 
# as published by the Free Software Foundation; either 
# version 2 of the License, or (at your option) any later 
# version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.


# ******************************************************************************
# ******************************************************************************
#
#
# CONSTANTS
#
#
# ******************************************************************************
# ******************************************************************************

USAGE=`cat << EOF

Usage: psnrcore ORIGINAL COMPARE PSNR_LOG [frames] [orig_opts] [comp_opts]

where ORIGINAL  is the path to the original video.
      COMPARE   is the path to the modified video (to be compared against 
                the original video file). Note: the same video may be used to
                test filters by using the *_opts parameters (see Example).
      PSNR_LOG  is a path (including the name) to a file to which the PSNR for 
                each frame will be written. If the log doesn't exist, it will
                be created.
      frames    is the number of frames to compare (starting from the first)
      orig_opts are the options to add to the mplayer command when creating
                the frames from the original video. "Surround in double quotes."
      comp_opts are the options to add to the mplayer command when creating
                the frames from the comparison video. "Use double quotes!"

Example:
           ./vidpsnr ./orig.avi ./orig.avi ./psnr-log.dat 250 "" "-vf pp=hb"
EOF`

ORIG_DIR=./orig
COMP_DIR=./comp

MPLAYER_OPTS="-benchmark -nosound -noframedrop -noautosub"


# ******************************************************************************
# ******************************************************************************
#
#
# BASIC SET-UP
#
#
# ******************************************************************************
# ******************************************************************************

# Sanity checks
# Need at least 3 input arguments
if test $# -le 2; then
   echo "$USAGE"
   exit 1
fi

ORIGINAL="$1"
COMPARE="$2"
PSNR_LOG="$3"

# Do the videos exist?
if test -f "$ORIGINAL"; then
   :
else
   echo "Could not find orignal video file: $ORIGINAL. Exiting."
   exit 1
fi

if test -f "$COMPARE"; then
   :
else
   echo "Could not find comparison video file: $COMPARE. Exiting."
   exit 1
fi

LAST_FRAME=-1
if [ $# -ge 4 ]; then
	LAST_FRAME=$4
	echo
	echo "Processing the first $LAST_FRAME frames."
fi

if [ $# -ge 5 ]; then
	ORIG_OPTS="$5"
	echo "MPlayer options for $ORIGINAL: $ORIG_OPTS"
fi

if [ $# -ge 6 ]; then
	COMP_OPTS="$6"
	echo "MPlayer options for $COMPARE: $COMP_OPTS"
fi

mkdir $ORIG_DIR
mkdir $COMP_DIR

rm -f $ORIG_DIR/*.ppm
rm -f $COMP_DIR/*.ppm


# ******************************************************************************
# ******************************************************************************
#
#
# EXTRACT FRAMES FROM VIDEOS AND EVALUATE PSNR
#
#
# ******************************************************************************
# ******************************************************************************

# Create frames from the original video file
echo
echo "Making frames from $ORIGINAL"

if test $LAST_FRAME -ge 0; then
	CMD="mplayer $MPLAYER_OPTS $ORIG_OPTS -frames $LAST_FRAME -vo pnm:outdir=$ORIG_DIR $ORIGINAL > /dev/null"
else
	CMD="mplayer $MPLAYER_OPTS $ORIG_OPTS -vo pnm:outdir=$ORIG_DIR $ORIGINAL > /dev/null"
fi

echo $CMD
eval $CMD

# Create frames from the comparison video file
echo
echo "Making frames from $COMPARE"

if test $LAST_FRAME -ge 0 ; then
	CMD="mplayer $MPLAYER_OPTS $COMP_OPTS -frames $LAST_FRAME -vo pnm:outdir=$COMP_DIR $COMPARE > /dev/null"
else
	CMD="mplayer $MPLAYER_OPTS $COMP_OPTS -vo pnm:outdir=$COMP_DIR $COMPARE > /dev/null"
fi

echo $CMD
eval $CMD

# Calculate the PSNR for the two videos
echo
echo "Finding PSNR between $ORIGINAL and $COMPARE"

if PSNR=`psnrcore $ORIG_DIR $COMP_DIR $PSNR_LOG`; then
   echo
   echo "PSNR: $PSNR"
   echo "Frame report: $PSNR_LOG"
   rm -f $ORIG_DIR/*.ppm
   rm -f $COMP_DIR/*.ppm
   rmdir $ORIG_DIR
   rmdir $COMP_DIR
else
   echo "There is a problem with psnrcore! Could not find the PSNR."
   exit 1
fi

exit 0